"""
=======================
Visualization utilities
=======================

.. note::
    Try on `Colab <https://colab.research.google.com/github/pytorch/vision/blob/gh-pages/main/_generated_ipynb_notebooks/plot_visualization_utils.ipynb>`_
    or :ref:`go to the end <sphx_glr_download_auto_examples_others_plot_visualization_utils.py>` to download the full example code.

This example illustrates some of the utilities that torchvision offers for
visualizing images, bounding boxes, segmentation masks and keypoints.
"""

# sphinx_gallery_thumbnail_path = "../../gallery/assets/visualization_utils_thumbnail2.png"

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import make_grid
from torch.autograd import Variable
import torchvision.transforms.functional as F
import cv2
import numpy as np
import matplotlib.pyplot as plt
import time


plt.rcParams['interactive'] = "true"
plt.rcParams["savefig.bbox"] = 'tight'

def show(imgs):
    if not isinstance(imgs, list):
        imgs = [imgs]
    fig, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    for i, img in enumerate(imgs):
        img = img.detach()
        img = F.to_pil_image(img)
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

class Model(torch.nn.Module):
    def __init__(self) :
        super(Model, self).__init__()
        self.conv1 = torch.nn.Sequential(torch.nn.Conv2d(1, 64, 3, 1, 1),
                                         torch.nn.ReLU(),
                                         torch.nn.Conv2d(64, 128, 3, 1, 1),
                                         torch.nn.ReLU(),
                                         torch.nn.MaxPool2d(2, 2)) 
        self.dense = torch.nn.Sequential(torch.nn.Linear(14*14*128, 1024),
                                         torch.nn.ReLU(),
                                         torch.nn.Dropout(p = 0.5),
                                         torch.nn.Linear(1024, 10))
        
    def forward(self, x) :
        x = self.conv1(x)
        x = x.view(-1, 14*14*128)
        x = self.dense(x)
        return x

train_dataset = datasets.MNIST(root= 'data/', train = True,
                               transform = transforms.ToTensor(), download = True)
test_dataset = datasets.MNIST(root = 'data/', train = False, 
                               transform = transforms.ToTensor(), download = True)

train_loader = DataLoader(dataset = train_dataset, batch_size = 100, shuffle = True)
test_loader = DataLoader(dataset = test_dataset, batch_size= 100, shuffle = True)

images, lables = next(iter(train_loader))
img = make_grid(images, nrow = 10)
img = img.numpy().transpose(1, 2, 0)

# cv2.imshow('img', img)
# cv2.waitKey(0)

torch.serialization.add_safe_globals([Model])
device = torch.device('cpu')
model = Model().to(device)
model.load_state_dict(torch.load('parameter_first.pkl', weights_only=False)) #load

print("start test..., len(test_dataset):%d" % len(test_dataset))
start_time = time.time()
test_correct = 0
for data in test_loader:
    inputs, lables = data
    inputs, lables = Variable(inputs).cpu(), Variable(lables).cpu()
    outputs = model(inputs)
    _, id = torch.max(outputs.data, 1)
    test_correct += torch.sum(id == lables.data)
elapsed_time = time.time() - start_time
print("correct:%.3f%%, spend time:%f" % (100 * test_correct / len(test_dataset), elapsed_time))

